Currently I have only implemented to_graph and it works for comp, branch, cell, net, although I have not yet added tests. The functionality is very neat and makes looking at the networks / cells really intuitive. For the following network,

comp = jx.Compartment()
branch = jx.Branch([comp for _ in range(4)])
cell = jx.Cell([branch for _ in range(5)], parents=jnp.asarray([-1, 0, 1, 2, 2]))
net = jx.Network([cell]*3)
connect(net[0,0,0], net[1,0,0], IonotropicSynapse())
connect(net[0,0,1], net[1,0,1], IonotropicSynapse())
connect(net[0,0,1], net[1,0,1], TestSynapse())
net.cell(2).add_to_group("cell2")
net.cell(2).branch(1).add_to_group("cell2brach1")


net.cell(0).insert(Na())
net.cell(0).insert(Leak())

net.cell(1).branch(1).insert(Na())
net.cell(0).insert(K())
net.compute_xyz()

net.cell(0).branch(0).loc(0.0).record()
net.cell(0).branch(0).loc(0.0).record("m")
current = jx.step_current(i_delay, i_dur, i_amp, dt, t_max)
net.cell(0).branch(2).loc(0.0).stimulate(current)
net.cell(0).branch(1).make_trainable("Na")
net.cell(1).make_trainable("K")

net.compute_xyz()
net.cell(1).move(0,30,0)
net.cell(2).move(0,-30,0)

we can just to_graph, plot it

module_graph = net.to_graph()
# plot the graph
pos = {i: (n["x"], n["y"]) for i, n in module_graph.nodes(data=True)} 
plt.figure(figsize=(8, 8))
nx.draw(module_graph, pos, with_labels=True, node_size=200, node_color="skyblue", font_size=8, font_weight="bold", font_color="black", font_family="sans-serif")
plt.show()

and look at all its properties, i.e. checking out the soma (0) and the synapse going from node 0 to node 20

print(module_graph.nodes[0])
print(module_graph.edges[(0,20)])

{'comp_index': 0, 'branch_index': 0, 'cell_index': 0, 'length': 10.0, 'radius': 1.0, 'axial_resistivity': 5000.0, 'capacitance': 1.0, 'v': -70.0, 'Na': True, 'Na_gNa': 0.05, 'Na_eNa': 35.0, 'Na_m': 0.2, 'Na_h': 0.2, 'Leak': True, 'Leak_gLeak': 5e-05, 'Leak_eLeak': -67.0, 'K': True, 'K_gK': 0.012, 'eK': -75.0, 'K_n': 0.1, 'x': 5.0, 'y': 0.0, 'z': 0.0, 'recordings': array(['v', 'm'], dtype=object)}

{'pre_locs': 0.125, 'pre_branch_index': 0, 'pre_cell_index': 0, 'post_locs': 0.125, 'post_branch_index': 0, 'post_cell_index': 1, 'type': 'IonotropicSynapse', 'type_ind': 0, 'global_pre_comp_index': 0, 'global_post_comp_index': 20, 'global_pre_branch_index': 0, 'global_post_branch_index': 5, 'IonotropicSynapse_gS': 0.5, 'IonotropicSynapse_e_syn': 0.0, 'IonotropicSynapse_k_minus': 0.025, 'IonotropicSynapse_s': 0.2, 'TestSynapse_gC': nan, 'TestSynapse_c': nan}

Both to_graph and from_graph work. If you want to give this a preliminary review, that would be awesome! Please also chime in if you have specific ideas about the format in which to best store the attributes in the graph. I think we do not have to stick with how it is stored in Module and should rather ensure that they are easy to parse and work with.

Yes, as mentioned in the PR desc, reading swc to graph is the plan.

Potential use-case of to_graph:

• reordering parents to have a different "initial" parent node
• search operations such as shortest path

It took way more time than I would have suspected or liked to spend and two mayor iterations, but its finally mostly there! Both export and import from/to graph are working now and also the swc pipeline passes the tests now. Some more tests and cleanup, but then I think you can do a review @michaeldeistler. Small benefit of this PR is also that it reduces the difference of jaxley and NEURON very slightly (by about 20%).

TLDR: I think I have a working version of the graph pipeline now, but tests are still not passing and I don't quite no why. Everything looks good to me. Help is much appreciated!

I have been trying for ages to get from_graph working. For this I have written a networkx -> jaxley converter. I think it works quite well. Running:

graph = swc_to_graph(fname)
cell = from_graph(graph, nseg=8, max_branch_len=2000.0)

reads an swc file into a networkx graph and imports the cell into jaxley. As part of from_graph the afformentioned converter is run, but it can also be run seperately via: graph = make_jaxley_compatible(graph). This produces a jx.Cell object that runs and does things. Additionally, inspecting the morphologies I get visually and comparing the compartment coordinates to NEURON / the current swc pipeline also looks good. As a bonus the compartment centers assigned via the graph are even closer to NEURON than the current ones. See the left plot.

While this all looks very promising, I have struggled to get the tests passing and I really dont know why:
For this I have replaced the hardcoded branch indices and pathlength matching in the old tests and I now just match the closest compartment to NEURON. This works quite well and I can even assign trunk_ids etc. based on the swc ids which is nice. The test should be much more robust to changes in the test_morphology or other things now! Of course it also passes for the old swc pipeline. However, when I run this test on the graph imported jx.Cell the voltages are consistently. See here:

While they look very similar, they are all slightly delayed. The MSEs also show this.

MSEs between networkX and NEURON voltages [10.14612018 10.35377231 10.84316607 11.17640196 11.06866644 11.1941421
 10.72721014 10.44691712 10.55542459  9.83932744  9.47578905 10.63712368
 13.2997599   9.5394632  10.3148209  10.06744197 10.11404793  9.99624842
 11.15663607 11.10845196 10.55436891 10.17346833 10.95011269 11.09521677
 11.24092038 10.65548104 11.36646986]
MSEs between jaxley_swc_reader and NEURON voltages [0.72412308 0.7978693  0.81918308 0.84724678 0.80464105 1.43258348
 1.47786316 0.73113843 1.14953477 1.12154859 1.18076938 1.30564702
 1.56853646 2.9105652  0.38078612 1.00738273 0.79437486 0.81636714
 0.48941009 0.47586723 0.75917783 1.04822786 1.35910234 1.56591299
 1.14104993 1.34106949 1.72583906]

Would really appreciate help. I will now be moving on to sth else for the moment though, so it is not urgent.
@michaeldeistler

Best,

Jonas

I think this is finally ready for a first round of reviews.

This has become quite the mammoth PR, but the functionality it enables is neat imo. For a rundown see the updated 08_morphologies.ipynb

All tests are passing now, which turned out to be an immense amount of work, but the imported morphologies are similar enough to NEURON both at the compartment level (x,y,z,r,l) and they also simulate correctly. I have essentially cloned the tests in tests/test_swc.py and added them to tests/test_graph.py. I have kept the tests and methods separated from everything else (although a refactor / integration might make it easier to maintain later). The import and export functionality also works correctly.

Notable changes are:

• new io submodule

  • with swc_to_graph, from_graph and to_graph functions
  • moved read_swc there from cell.py
  • added io.graph
  • fixed the _update_nodes_with_xyz since it had the same issue that caused bug in compute_xyz #410.
  • added __eq__ to Module (I needed this to test whether re-imported modules are the same as the original.)
  • test_graph.py matches voltages between NEURON and Jaxley by matching compartment centers, this makes the test much more robust. Could also be added to test_swc.py if desired (see tests/helpers.py).

• Some open questions are:

  • __eq__ is only implicitely being tested, and I dont know whether there could be a more efficient implementation
  • there is a lot of overlap between test_graph and test_swc. We could potentially move them to tests/io/ and reduce overlap.

Lemme know your thoughts. Would also be happy to go through this in person.

Just noting that in the branch graph_distance_and_more, there is some preliminary API idea that works, e.g., with, e.g., "distance" as a new function. Implementation may be inefficient yet, but can be improved if we can also add the graph to the module...

It is a branch from this branch at the current stage.

This has become a marathon PR, due to what it took to get the tests to pass (matching NEURON's and jaxley's morphologies by coordinates, simulating swc errors, ... just to name a few). However, I think it is on the home stretch. All tests are passing and I think all functionality is there. @michaeldeistler I left you a few comments. Your feedback would be greatly appreciated.

Two things:

1. @manuelgloeckler distances should now be easy to make work.
2. Should we merge this with main or v1.0

Good to know its robust to new morphologies :D I will add this to the testcases.

I have to admit that I did not do a detailed review of the rest of the code---but given that I love the API and the tutorial and all tests (including complicated ones as described above) pass, we should merge this.

I can give you a rundown when your back next week. Before I merge this though, @michaeldeistler could you just briefly go through the points in my own review above? Would appreciate if you could leave your opinion.

One remaining question: do I understand correctly that read_swc does the same thing as before? IIUC, we do not yet have any SWC->networkX function, right? One has to go via the Jaxley swc reader?

I kept the old pipeline completely intact (incl. read_swc), but we can think about getting rid of it longterm if we want.
There is an swc->networkX function. See tutorial its called swc_to_graph and the only thing is does is read the swc file into a networkx graph. Also have a look here. You can in principle obtain a graph at various stages of conversion, i.e. with incl branches, incl comps, or just the comps.

There is also a lot of context in the docstrings. :)

@michaeldeistler marked the test to be skipped for now. Lemme know if you're ok with merging this.